import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import rcParams

# 切换到支持 Unicode 的字体
rcParams['font.family'] = 'Arial Unicode MS'

# 读取数据
df = pd.read_csv('data/merged_olympics_data.csv')
df.columns = df.columns.str.strip().str.upper()

# 计算金、银、铜牌
df['GOLD'] = (df['MEDAL'] == 'Gold').astype(int)
df['SILVER'] = (df['MEDAL'] == 'Silver').astype(int)
df['BRONZE'] = (df['MEDAL'] == 'Bronze').astype(int)
df['TOTAL'] = df['GOLD'] + df['SILVER'] + df['BRONZE']

# 事件清理
df['EVENT'] = df['EVENT'].str.strip()

# 聚合按国家和事件
country_event_summary = df.groupby(['NOC', 'EVENT']).agg({
    'GOLD': 'sum',
    'SILVER': 'sum',
    'BRONZE': 'sum',
    'TOTAL': 'sum'
}).reset_index()

# 计算比例和分数
country_event_summary['GOLD_RATIO'] = country_event_summary['GOLD'] / country_event_summary['TOTAL']
country_event_summary['SILVER_RATIO'] = country_event_summary['SILVER'] / country_event_summary['TOTAL']
country_event_summary['BRONZE_RATIO'] = country_event_summary['BRONZE'] / country_event_summary['TOTAL']

country_event_summary['GOLD_SCORE'] = country_event_summary['GOLD'] * country_event_summary['GOLD_RATIO']
country_event_summary['SILVER_SCORE'] = country_event_summary['SILVER'] * country_event_summary['SILVER_RATIO']
country_event_summary['BRONZE_SCORE'] = country_event_summary['BRONZE'] * country_event_summary['BRONZE_RATIO']

# 创建 IMPORTANCE_SCORE 列
country_event_summary['IMPORTANCE_SCORE'] = country_event_summary[['GOLD_SCORE', 'SILVER_SCORE', 'BRONZE_SCORE']].sum(axis=1)

# 过滤国家：只保留中国、美国、日本
filtered_data = country_event_summary[country_event_summary['NOC'].isin(['CHN', 'USA', 'JPN'])]

# 定义函数获取前10事件
def get_top_events(data, country):
    country_data = data[data['NOC'] == country].sort_values(by='IMPORTANCE_SCORE', ascending=False).head(10)
    return country_data

# 获取每个国家前10的重要事件
chn_top = get_top_events(filtered_data, 'CHN')
usa_top = get_top_events(filtered_data, 'USA')
jpn_top = get_top_events(filtered_data, 'JPN')

# 绘图
plt.figure(figsize=(14, 7))

for country_data, label, color in zip(
    [chn_top, usa_top, jpn_top],
    ['China (CHN)', 'USA', 'Japan (JPN)'],
    ['red', 'blue', 'orange']  # 修改日本的颜色为橙色
):
    plt.scatter(
        country_data['EVENT'],
        country_data['IMPORTANCE_SCORE'],
        label=label,
        color=color,
        s=100,  # 点的大小
        alpha=0.8,  # 透明度
        edgecolors='black'  # 添加点边框
    )

plt.title('Top 10 Important Events and Scores for China, USA, and Japan', fontsize=16)
plt.xlabel('Event', fontsize=12)
plt.ylabel('Importance Score', fontsize=12)
plt.xticks(rotation=45, ha='right')
plt.legend(title='Country', fontsize=10)
plt.grid(axis='y', linestyle='--', alpha=0.7)
plt.tight_layout()

# 显示图像
plt.show()
